from pathlib import Path

from PIL import Image
from torch.utils.data import Dataset

from yms_zsl.dataset_utils.FactorLabelMapper import FactorLabelMapper


class ImageClassificationDataset(Dataset):
    """
    自定义数据集类
    """

    def __init__(self, root_dir, transform=None, train_class=None, mode='train', ignore_factors=None,
                 factor_index_map_path=None):
        """
        初始化数据集

        参数:
            root_dir (str): 数据集根目录路径
            transform (callable, optional): 应用于图像的变换/增强
            target_transform (callable, optional): 应用于标签的变换
            mode (str): 数据集模式，如 'train', 'val', 'test'
        """
        self.root_dir = Path(root_dir)
        self.mode = mode
        self.transform = transform

        # 构建数据路径
        self.data_dir = self.root_dir / mode

        # 确保目录存在
        if not self.data_dir.exists():
            raise ValueError(f"目录不存在: {self.data_dir}")

        # 2. 初始化训练类别映射器
        self.mapper = FactorLabelMapper(
            data_dir=self.data_dir,  # 当train_class不存在时，从当前模式目录提取训练类别
            class_list_path=train_class,
            ignore_factors=ignore_factors,
            factor_index_map_path=factor_index_map_path
        )
        self.train_classes = self.mapper.raw_classes  # 训练使用的类别
        self.classes = self.mapper.classes  # 实际预测的所有类别
        self.class_to_idx = self.mapper.class_to_idx  # 基于预测类别构建的映射

        # 收集所有图像路径和对应的标签
        self.image_paths = []
        self.labels = []

        for class_name in self.train_classes:
            class_dir = self.data_dir / class_name
            class_idx = self.class_to_idx[class_name]

            # 获取此类别的所有图像文件
            image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
            for ext in image_extensions:
                for img_path in class_dir.glob(f'*{ext}'):
                    self.image_paths.append(img_path)
                    self.labels.append(class_idx)

        if not self.image_paths:
            raise ValueError(f"在 {self.data_dir} 中未找到任何图像文件")

    def __len__(self):
        """返回数据集中的样本数量"""
        return len(self.image_paths)

    def __getitem__(self, idx):
        """
        获取单个样本

        参数:
            idx (int): 样本索引

        返回:
            tuple: (图像, 标签)
        """
        img_path = self.image_paths[idx]
        label = self.labels[idx]

        # 加载图像
        image = Image.open(img_path).convert('RGB')  # 确保图像是RGB格式

        # 应用变换
        if self.transform:
            image = self.transform(image)

        return image, label

    def get_class_distribution(self):
        """返回每个类别的样本数量"""
        class_counts = {cls: 0 for cls in self.classes}
        for label in self.labels:
            class_name = self.classes[label]
            class_counts[class_name] += 1
        return class_counts

    def get_class_num(self):
        return len(self.classes)


class FeatureDecouplingDataset(Dataset):
    def __init__(self, root_dir, transform=None, train_class=None, mode='train', factor_index_map_path=None, ignore_factors=None):
        self.root_dir = Path(root_dir)
        self.transform = transform
        self.data_dir = self.root_dir / mode

        self.maper = FactorLabelMapper(self.data_dir, train_class,
                                       factor_index_map_path=factor_index_map_path, ignore_factors=ignore_factors)
        self.train_classes = self.maper.raw_classes

        self.classes = self.maper.classes

        # 创建类别到索引的映射
        self.class_to_idx = self.maper.class_to_idx

        # 收集所有图像路径和对应的标签
        self.image_paths = []
        self.labels = []
        self.indices = []

        for class_name in self.train_classes:
            class_dir = self.data_dir / class_name
            class_idx = self.class_to_idx[class_name]

            # 获取此类别的所有图像文件
            image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
            for ext in image_extensions:
                for img_path in class_dir.glob(f'*{ext}'):
                    self.image_paths.append(img_path)
                    self.labels.append(class_idx)
                    self.indices.append(
                        self.maper.get_indices_from_factors(self.maper.get_factors_from_class(class_name)))

        if not self.image_paths:
            raise ValueError(f"在 {self.data_dir} 中未找到任何图像文件")

    def __len__(self):
        return len(self.image_paths)

    def __getitem__(self, idx):
        img_path = self.image_paths[idx]
        label = self.labels[idx]

        # 加载图像
        image = Image.open(img_path).convert('RGB')  # 确保图像是RGB格式

        # 应用变换
        if self.transform:
            image = self.transform(image)
        indices = self.indices[idx]

        return image, indices, label


class ImageDataset(Dataset):
    """
    自定义图像分类数据集类（直接从data_dir子文件夹获取类别，无FactorLabelMapper依赖）
    支持标准数据集结构：root_dir/mode/类别文件夹/图像文件
    """

    def __init__(self, root_dir, transform=None, mode='train'):
        """
        初始化数据集

        参数:
            root_dir (str): 数据集根目录路径（根目录下包含mode对应的子文件夹）
            transform (callable, optional): 应用于图像的变换/增强（如Resize、ToTensor等）
            target_transform (callable, optional): 应用于标签的变换（如OneHotEncoder等）
            mode (str): 数据集模式，对应根目录下的子文件夹名（如 'train', 'val', 'test'）
        """
        self.root_dir = Path(root_dir)
        self.mode = mode
        self.transform = transform

        # 1. 构建当前模式的数据集目录（root_dir/mode，如"数据根目录/train"）
        self.data_dir = self.root_dir / mode
        # 检查目录是否存在
        if not self.data_dir.exists():
            raise ValueError(f"当前模式的数据集目录不存在: {self.data_dir}")

        # 2. 直接从data_dir获取类别（子文件夹名称即为类别名）
        # 过滤：仅保留子文件夹，排除文件或隐藏目录（如.git、.DS_Store等）
        self.classes = [
            item.name for item in self.data_dir.iterdir()
            if item.is_dir() and not item.name.startswith('.')
        ]
        # 检查是否有类别文件夹
        if not self.classes:
            raise ValueError(f"在 {self.data_dir} 中未找到任何类别子文件夹，请检查数据集结构")
        # 对类别排序（保证类别→索引映射的一致性，避免不同运行顺序导致索引混乱）
        self.classes.sort()

        # 3. 构建类别→索引的映射（按排序后的类别顺序分配索引）
        self.class_to_idx = {cls_name: idx for idx, cls_name in enumerate(self.classes)}

        # 4. 收集所有图像路径和对应的标签
        self.image_paths = []
        self.labels = []
        # 支持的图像文件后缀（覆盖常见格式）
        self.image_extensions = ('.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.gif', '.ppm')

        # 遍历每个类别文件夹，收集图像
        for cls_name in self.classes:
            cls_dir = self.data_dir / cls_name  # 类别文件夹路径
            cls_idx = self.class_to_idx[cls_name]  # 类别对应的索引

            # 遍历类别文件夹下所有符合后缀的图像文件
            for img_path in cls_dir.glob(f'*'):  # 先获取所有文件
                if img_path.suffix.lower() in self.image_extensions:  # 过滤图像后缀
                    self.image_paths.append(img_path)
                    self.labels.append(cls_idx)

        # 检查是否收集到图像
        if not self.image_paths:
            raise ValueError(f"在 {self.data_dir} 下的类别文件夹中未找到任何图像文件（支持格式：{self.image_extensions}）")

    def __len__(self):
        """返回数据集中的样本总数"""
        return len(self.image_paths)

    def __getitem__(self, idx):
        """
        根据索引获取单个样本（图像+标签）

        参数:
            idx (int): 样本索引

        返回:
            tuple: (处理后的图像张量, 处理后的标签)
        """
        # 获取图像路径和原始标签
        img_path = self.image_paths[idx]
        label = self.labels[idx]

        image = Image.open(img_path).convert('RGB')

        # 应用图像变换（如Resize、ToTensor、Normalize等）
        if self.transform:
            image = self.transform(image)

        return image, label

    def get_class_distribution(self):
        """返回每个类别的样本数量（用于分析数据分布）"""
        class_counts = {cls: 0 for cls in self.classes}
        for label_idx in self.labels:
            cls_name = self.classes[label_idx]  # 从索引反查类别名
            class_counts[cls_name] += 1
        return class_counts

    def get_class_num(self):
        """返回数据集中的类别总数"""
        return len(self.classes)

    def get_class_mapping(self):
        """返回类别→索引的映射字典（用于后续结果解析，如索引转类别名）"""
        return self.class_to_idx.copy()
